'''关于路径的常用处理代码.'''

import os
from contextlib import contextmanager
from note_utils.typing import Any, Generator, StrPath
from pathlib import Path

def expand_path(path:StrPath) -> Path:
    '''拓展路径中的家目录和环境变量,返回绝对路径.'''
    path=os.path.expandvars(path)
    path=os.path.expanduser(path)
    path=os.path.abspath(path)
    return Path(path)

@contextmanager
def chdir(
    path: StrPath
) -> Generator['tuple[Path,Path]', Any,
               None]:  # for py<3.11,slightly different
    '''暂时切换到指定目录,with块结束时回到切换前的目录.
    
    生成的绑定对象是(指定目录,原始目录)的绝对路径.
    
    >>> import os
    >>> os.chdir('/etc')
    >>> os.getcwd()
    '/etc'
    >>> with chdir('/run') as (tmp,origin):
    ...     print(f'with块中的当前目录:{tmp}')
    with块中的当前目录:/run
    >>> print(f'当前目录:{os.getcwd()}')
    当前目录:/etc
    '''
    path_init = Path(os.getcwd())
    tmp_path = expand_path(path)
    try:
        os.chdir(tmp_path)
        print(f'cd {tmp_path.as_posix()}')
        print(f'当前目录从 {path_init.as_posix()} \n临时切换到 {tmp_path.as_posix()}')
        yield (tmp_path, path_init)
    finally:
        os.chdir(path_init)
        print(f'当前目录切换回到 {path_init.as_posix()}')


def add_rosetta_env(rosetta_home: StrPath) -> Path:
    '''添加ROSETTAHOME,ROSETTA3,ROSETTA3,ROSETTA3_DB,ROSETTA_TOOLS环境变量到笔记本中.
    
    返回rosetta_home的Path对象.'''
    rosetta_home=expand_path(rosetta_home)
    os.environ['ROSETTAHOME']=str(rosetta_home)
    os.environ['ROSETTA3']=f'{rosetta_home}/source'
    os.environ['ROSETTA3_DB']=f'{rosetta_home}/database'
    os.environ['ROSETTA_TOOLS']=f'{rosetta_home}/tools'
    return rosetta_home
